import logging
from datetime import datetime
import json


class JsonFormatter(logging.Formatter):
    
    def __init__(
        self,
        fmt=None, datefmt=None, style='%', validate=True,
        fields: list = None,
    ):
        super().__init__(fmt, datefmt, style, validate)
        self.fields = [] if '*' in fields or not fields else fields
    
    BASE_NEED_ATTR = ["levelname", "asctime", "module", "funcName", "created", "process", "thread", "filename", "lineno", "levelno", "created", "exc_info", "message"]

    def format(self, record):
        extra = self.build_record(record)
        filter_fields = set(self.fields) & set(self.BASE_NEED_ATTR) if self.fields else set(self.BASE_NEED_ATTR)
        extra = {key: extra[key] for key in extra if key in filter_fields}
        self.set_format_time(extra)  # set time
        extra['message'] = record.msg if record.message else record.msg
        if record.exc_info:
            extra['exc_info'] = self.formatException(record.exc_info)
        if self._fmt == 'pretty':
            return json.dumps(extra, indent=1, ensure_ascii=False)
        else:
            return json.dumps(extra, ensure_ascii=False)

    @classmethod
    def build_record(cls, record):
        return {
            attr_name: record.__dict__[attr_name]
            for attr_name in record.__dict__
            if attr_name in cls.BASE_NEED_ATTR
        }

    @classmethod
    def set_format_time(cls, extra):
        now = datetime.utcnow()
        format_time = now.strftime("%Y-%m-%dT%H:%M:%S" + ".%03d" % (now.microsecond / 1000) + "Z")
        extra['@timestamp'] = format_time
        return format_time